package com.lijiajia.cloud.authserver.security.oauth.token;

import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.security.oauth2.common.ExpiringOAuth2RefreshToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.DefaultAuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.TokenStore;

import java.time.OffsetDateTime;
import java.util.Collection;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * redis token store
 *
 * @author lijiajia
 * @since 2019-3-29
 */
@Slf4j
public class PlusTokenStore implements TokenStore {

    private final RedisTemplate<String, Object> redisTemplate;

    private static final String ACCESS = "access:";
    private static final String AUTH_TO_ACCESS = "auth_to_access:";
    private static final String AUTH = "auth:";
    private static final String REFRESH_AUTH = "refresh_auth:";
    private static final String ACCESS_TO_REFRESH = "access_to_refresh:";
    private static final String REFRESH = "refresh:";
    private static final String REFRESH_TO_ACCESS = "refresh_to_access:";
    private static final String CLIENT_ID_TO_ACCESS = "client_id_to_access:";
    private static final String USERNAME_TO_ACCESS = "username_to_access:";

    private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();

    public PlusTokenStore(RedisTemplate<String, Object> redisTemplate) {
        log.debug("token plus 配置");
        this.redisTemplate = redisTemplate;
    }

    /**
     * 读取存储在指定标记值下的身份验证。
     *
     * @param token 存储身份验证的令牌值。
     * @return 身份验证，如果没有，则返回null。
     */
    @Override
    public OAuth2Authentication readAuthentication(OAuth2AccessToken token) {
        return readAuthentication(token.getValue());
    }

    /**
     * 读取存储在指定标记值下的auth。
     *
     * @param token 存储auth的token。
     * @return 身份验证，如果没有，则返回null
     */
    @Override
    public OAuth2Authentication readAuthentication(String token) {
        return (OAuth2Authentication) this.redisTemplate.opsForValue().get(AUTH + token);
    }

    /**
     * 存储访问token
     *
     * @param token          要存储的令牌
     * @param authentication 与token关联的auth
     */
    @Override
    public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {
        this.redisTemplate.opsForValue().set(ACCESS + token.getValue(), token);
        this.redisTemplate.opsForValue().set(AUTH + token.getValue(), authentication);
        this.redisTemplate.opsForValue().set(AUTH_TO_ACCESS + authenticationKeyGenerator.extractKey(authentication), token);
        if (!authentication.isClientOnly()) {
            redisTemplate.opsForList().rightPush(USERNAME_TO_ACCESS + getApprovalKey(authentication), token);
        }

        redisTemplate.opsForList().rightPush(CLIENT_ID_TO_ACCESS + authentication.getOAuth2Request().getClientId(), token);

        if (token.getExpiration() != null) {
            int seconds = token.getExpiresIn();
            redisTemplate.expire(ACCESS + token.getValue(), seconds, TimeUnit.SECONDS);
            redisTemplate.expire(AUTH + token.getValue(), seconds, TimeUnit.SECONDS);
            redisTemplate.expire(AUTH_TO_ACCESS + authenticationKeyGenerator.extractKey(authentication), seconds, TimeUnit.SECONDS);
            redisTemplate.expire(CLIENT_ID_TO_ACCESS + authentication.getOAuth2Request().getClientId(), seconds, TimeUnit.SECONDS);
            redisTemplate.expire(USERNAME_TO_ACCESS + getApprovalKey(authentication), seconds, TimeUnit.SECONDS);
        }
        if (token.getRefreshToken() != null && token.getRefreshToken().getValue() != null) {
            this.redisTemplate.opsForValue().set(REFRESH_TO_ACCESS + token.getRefreshToken().getValue(), token.getValue());
            this.redisTemplate.opsForValue().set(ACCESS_TO_REFRESH + token.getValue(), token.getRefreshToken().getValue());
            OAuth2RefreshToken refreshToken = token.getRefreshToken();
            if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
                ExpiringOAuth2RefreshToken expiringRefreshToken = (ExpiringOAuth2RefreshToken) refreshToken;
                Date expiration = expiringRefreshToken.getExpiration();
                if (expiration != null) {
                    long seconds = expiration.getTime() - OffsetDateTime.now().toEpochSecond();
                    redisTemplate.expire(REFRESH_TO_ACCESS + token.getRefreshToken().getValue(), seconds, TimeUnit.SECONDS);
                    redisTemplate.expire(ACCESS_TO_REFRESH + token.getValue(), seconds, TimeUnit.SECONDS);
                }
            }
        }
    }

    /**
     * 从商店中读取访问令牌。
     *
     * @param tokenValue token值
     * @return 要读取的token令牌
     */
    @Override
    public OAuth2AccessToken readAccessToken(String tokenValue) {
        return (OAuth2AccessToken) redisTemplate.opsForValue().get(ACCESS + tokenValue);
    }

    /**
     * 从商店中删除访问令牌
     *
     * @param token 要从商店中删除的令牌
     */
    @Override
    public void removeAccessToken(OAuth2AccessToken token) {
        removeAccessToken(token.getValue());
    }

    public void removeAccessToken(String tokenValue) {
        OAuth2Authentication authentication = (OAuth2Authentication) this.redisTemplate.opsForValue().get(AUTH + tokenValue);

        redisTemplate.delete(ACCESS + tokenValue);
        redisTemplate.delete(AUTH + tokenValue);
        redisTemplate.delete(ACCESS_TO_REFRESH + tokenValue);
        if (authentication != null) {
            this.redisTemplate.delete(AUTH_TO_ACCESS + authenticationKeyGenerator.extractKey(authentication));

            String clientId = authentication.getOAuth2Request().getClientId();
            redisTemplate.opsForList().leftPop(USERNAME_TO_ACCESS + getApprovalKey(clientId, authentication.getName()));
            redisTemplate.opsForList().leftPop(CLIENT_ID_TO_ACCESS + clientId);

            this.redisTemplate.delete(AUTH_TO_ACCESS + authenticationKeyGenerator.extractKey(authentication));
        }
    }

    /**
     * 将指定的刷新令牌存储在商店中。
     *
     * @param refreshToken   要存储的刷新令牌。
     * @param authentication 与刷新令牌关联的身份验证。
     */
    @Override
    public void storeRefreshToken(OAuth2RefreshToken refreshToken, OAuth2Authentication authentication) {
        redisTemplate.opsForValue().set(REFRESH + refreshToken.getValue(), refreshToken);
        redisTemplate.opsForValue().set(REFRESH_AUTH + refreshToken.getValue(), authentication);
        if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
            ExpiringOAuth2RefreshToken expiringRefreshToken = (ExpiringOAuth2RefreshToken) refreshToken;
            Date expiration = expiringRefreshToken.getExpiration();
            if (expiration != null) {
                long seconds = expiration.getTime() - OffsetDateTime.now().toEpochSecond();
                redisTemplate.expire(REFRESH + refreshToken.getValue(), seconds, TimeUnit.SECONDS);
                redisTemplate.expire(REFRESH_AUTH + refreshToken.getValue(), seconds, TimeUnit.SECONDS);
            }
        }
    }

    /**
     * 从商店中读取刷新令牌
     *
     * @param tokenValue 要读取的令牌的值
     * @return 令牌
     */
    @Override
    public OAuth2RefreshToken readRefreshToken(String tokenValue) {
        return (OAuth2RefreshToken) redisTemplate.opsForValue().get(REFRESH + tokenValue);
    }

    /**
     * @param token 刷新令牌
     * @return 最初用于授予刷新令牌的身份验证
     */
    @Override
    public OAuth2Authentication readAuthenticationForRefreshToken(OAuth2RefreshToken token) {
        return readAuthenticationForRefreshToken(token.getValue());
    }

    public OAuth2Authentication readAuthenticationForRefreshToken(String token) {
        return (OAuth2Authentication) redisTemplate.opsForValue().get(REFRESH_AUTH + token);
    }

    /**
     * 从商店中删除刷新令牌。
     *
     * @param refreshToken The token to remove from the store.
     */
    @Override
    public void removeRefreshToken(OAuth2RefreshToken refreshToken) {
        removeRefreshToken(refreshToken.getValue());
    }

    public void removeRefreshToken(String tokenValue) {
        redisTemplate.delete(REFRESH + tokenValue);
        redisTemplate.delete(REFRESH_AUTH + tokenValue);
        redisTemplate.delete(REFRESH_TO_ACCESS + tokenValue);
        redisTemplate.delete(ACCESS_TO_REFRESH + tokenValue);
    }

    /**
     * 使用刷新令牌删除访问令牌。此功能是必需的，因此无法使用刷新令牌
     * 创建无限数量的访问令牌
     *
     * @param refreshToken 刷新令牌
     */
    @Override
    public void removeAccessTokenUsingRefreshToken(OAuth2RefreshToken refreshToken) {
        removeAccessTokenUsingRefreshToken(refreshToken.getValue());
    }

    private void removeAccessTokenUsingRefreshToken(String refreshTokenValue) {
        redisTemplate.delete(REFRESH_TO_ACCESS + refreshTokenValue);
    }

    /**
     * 检索针对提供的身份验证密钥存储的访问令牌（如果存在）。
     *
     * @param authentication 访问令牌的身份验证密钥
     * @return 访问令牌，如果没有则为null
     */
    @Override
    public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
        String key = authenticationKeyGenerator.extractKey(authentication);
        OAuth2AccessToken accessToken = (OAuth2AccessToken) redisTemplate.opsForValue().get(AUTH_TO_ACCESS + key);
        if (accessToken != null) {
            OAuth2Authentication storedAuthentication = readAuthentication(accessToken.getValue());
            // 保持多点登陆用户信息一致
            if (storedAuthentication == null || !key.equals(authenticationKeyGenerator.extractKey(storedAuthentication))) {
                storeAccessToken(accessToken, authentication);
            }
        }
        return accessToken;
    }

    /**
     * @param clientId clientId要搜索的客户端ID
     * @param userName userName要搜索的用户名
     * @return 访问令牌的集合
     */
    @Override
    public Collection<OAuth2AccessToken> findTokensByClientIdAndUserName(String clientId, String userName) {
        String key = getApprovalKey(clientId, userName);
        List<Object> list = redisTemplate.opsForList().range(key, 1, -1);
        if (list == null || list.isEmpty()) {
            return Collections.emptySet();
        }
        return list.stream()
                .map(accessToken -> (OAuth2AccessToken) accessToken)
                .collect(Collectors.toUnmodifiableList());
    }

    /**
     * @param clientId 要搜索的客户端ID
     * @return 访问令牌的集合
     */
    @Override
    public Collection<OAuth2AccessToken> findTokensByClientId(String clientId) {
        List<Object> list = redisTemplate.opsForList().range(CLIENT_ID_TO_ACCESS + clientId, 0, -1);
        if (list == null || list.isEmpty()) {
            return Collections.emptySet();
        }
        return list.stream()
                .map(accessToken -> (OAuth2AccessToken) accessToken)
                .collect(Collectors.toUnmodifiableList());
    }

    private static String getApprovalKey(OAuth2Authentication authentication) {
        String userName = authentication.getUserAuthentication() == null ? ""
                : authentication.getUserAuthentication().getName();
        return getApprovalKey(authentication.getOAuth2Request().getClientId(), userName);
    }

    private static String getApprovalKey(String clientId, String userName) {
        return clientId + (userName == null ? "" : ":" + userName);
    }
}
